import torch
import torch.utils.data as Data

BATCH_SIZE = 6  # 批训练的数据个数

# x = torch.linspace(1, 100, 100)  # x data (torch tensor)
# y = torch.linspace(100, 1, 100)  # y data (torch tensor)
x = torch.rand(5, 5)  # x data (torch tensor)
y = torch.rand(5, 5)  # y data (torch tensor)

# 先转换成 torch 能识别的 Dataset
torch_dataset = Data.TensorDataset(x, y)

# 把 dataset 放入 DataLoader  Dataloader 模块快速导入信息
loader = Data.DataLoader(
    dataset=torch_dataset,  # torch TensorDataset format
    batch_size=BATCH_SIZE,  # mini batch size
    shuffle=True,  # 要不要打乱数据 (打乱比较好)
    # num_workers=2,  # 多线程来读数据
)

for epoch in range(3):  # 整体训练3次
    for step, (batch_x, batch_y) in enumerate(loader):
        # training.....
        print('Epoch: ', epoch, '| Step: ', step, '| batch x: ',
              batch_x.numpy(), '| batch y: ', batch_y.numpy())

